import cv2
import numpy as np
import matplotlib.pyplot as plt
import math
def main():
 grayimg = cv2.imread('camera.png', cv2.IMREAD_GRAYSCALE)
 plt.subplot(331), plt.imshow(grayimg, cmap='gray'), plt.title('original img'), plt.axis('off')
 grayimg = np.float64(grayimg/255.0)
 f = np.fft.fft2(grayimg)
 magnitude_spectrum1 = np.log(np.abs(f)+1)
 plt.subplot(332), plt.imshow(magnitude_spectrum1, cmap='gray'), plt.title('FFT before shift'), plt.axis('off')
 fshift = np.fft.fftshift(f)
 magnitude_spectrum2 = np.log(np.abs(fshift) + 1)
 plt.subplot(333), plt.imshow(magnitude_spectrum2, cmap='gray'), plt.title('FFT after shift'), plt.axis('off')
 rows, cols = grayimg.shape
 crow, ccol = int(rows / 2), int(cols / 2) # 找到中心位置
 # 理想低通滤波掩膜
 D0 = 50
 mask = np.zeros((rows, cols), np.float32)
 for i in range(rows):
  for j in range(cols):
   if (i - crow) ** 2 + (j - ccol) ** 2 <= D0 ** 2:
    mask[i, j] = 1.0
 # 理想高通滤波掩膜
 D1 = 50
 mask2 = np.ones((rows, cols), np.float32)
 for i in range(rows):
  for j in range(cols):
   if (i - crow) ** 2 + (j - ccol) ** 2 <= D1 ** 2:
    mask2[i, j] = 0.0
 m_fshift = fshift * mask
 m_fshift2 = fshift * mask2
 ishift = np.fft.ifftshift(m_fshift)
 ishift2 = np.fft.ifftshift(m_fshift2)
 iimg = np.fft.ifft2(ishift)
 iimg2 = np.fft.ifft2(ishift2)
 iimg = np.abs(iimg)
 iimg2 = np.abs(iimg2)
 plt.subplot(334), plt.imshow(mask, cmap='gray'), plt.title('low pass filter'), plt.axis('off')
 plt.subplot(337), plt.imshow(iimg, cmap='gray'), plt.title('low pass filterd image'), plt.axis('off')
 plt.subplot(335), plt.imshow(mask2, cmap='gray'), plt.title('high pass filter'), plt.axis('off')
 plt.subplot(338), plt.imshow(iimg2, cmap='gray'), plt.title('high pass filterd image'), plt.axis('off')
 plt.show()
if __name__ == '__main__':
 main()